import numpy as np

import pylidc as pl
from pylidc.utils import consensus

length = 0

# 遍历每一位患者
for patient_no in range(1, 1013):
    num = 'LIDC-IDRI-' + str(patient_no).zfill(4)
    print(num)

    # 获取该患者的所有扫描
    scans = pl.query(pl.Scan).filter(pl.Scan.patient_id == num)
    len_scan = scans.count()

    # 遍历每一次扫描
    for scan_no in range(len_scan):
        vol = scans[scan_no].to_volume()
        # 获取该次扫描下的所有结节标注
        nods = scans[scan_no].cluster_annotations()
        # 遍历每一个结节
        for nod_num in range(len(nods)):
            anns = nods[nod_num]

            # 只保留四个医生都进行标注的情况
            if len(anns) == 4:

                # consensus是用共识机制获取最优标注
                # 其中pad 三个tuple分别表示x,y,z轴，正负方向的拓展
                # 第一次consensus是为了获取共识标注的大小，方便后面加padding
                contour_mask, image_index, masks = consensus(anns, clevel=0.5,
                                                             pad=[(0, 0), (0, 0), (0, 0)])

                # 先获取标注的x,y,z轴大小
                x, y, z = map(int, contour_mask.shape)
                print(contour_mask.shape, x, y, z)
                # 填充至 (80，80，40)
                # 如果本身就有维度大于(80，80，40)，直接按照本身维度来
                x_pad = [40 if x < 80 else int(x / 2), 40 if x < 80 else int(x / 2 + 0.5)]
                y_pad = [40 if y < 80 else int(y / 2), 40 if y < 80 else int(y / 2 + 0.5)]
                z_pad = [20 if z < 40 else int(z / 2), 20 if z < 40 else int(z / 2 + 0.5)]
                # 这里需要一边int(x) 一边int(x+0.5),防止都被砍掉产生误差

                # 第二次consensus直接获取训练用的image和mask的index
                contour_mask, image_index, masks = consensus(anns, clevel=0.5,
                                                             pad=[(x_pad[0] - int(x / 2), x_pad[1] - int(x / 2 + 0.5)),
                                                                  (x_pad[0] - int(y / 2), y_pad[1] - int(y / 2 + 0.5)),
                                                                  (z_pad[0] - int(z / 2), z_pad[1] - int(z / 2 + 0.5))])

                # 截出图片和标注
                image = np.asarray(vol[image_index][:, :, :]).transpose(2, 0, 1)
                mask = np.float32(np.array(contour_mask[:, :, :])).transpose(2, 0, 1)

                if image.shape != (40, 80, 80) or mask.shape != (40, 80, 80):
                    print("error shape")
                    print(image.shape)
                    print(mask.shape)

                # 保存
                np.save("dataset/image/{}.npy".format(str(length)), image)
                np.save("dataset/mask/{}.npy".format(str(length)), mask)
                # np.savetxt("dataset/image/{}.csv".format(str(length)), image[0], delimiter=",")
                # np.savetxt("dataset/mask/{}.csv".format(str(length)), mask[0], delimiter=",")
                length += 1

    print(length)
    #  897
